import os

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import networkx as nx
from General.ConfigLoader import load_config
from General.DatabaseReader import DatabaseReader


class MapLoader:
    """
        地图加载类，将医院地图建模为点位与路段组成的图结构
    """
    def __init__(self):
        config = load_config()
        current_dir = os.path.dirname(os.path.abspath(__file__))
        self.code_index_path = os.path.join(current_dir, config['filepath']['node_code_index_map'])
        self.adjacency_matrix = None
        self.graph = None
        self.node_features = None
        self.edge_features = None
        self.code_index_map = None
        self.database_reader = DatabaseReader()
        self.load_node_code_index_map()

    def load_node_code_index_map(self):
        """加载节点代码索引映射"""
        self.code_index_map = pd.read_excel(self.code_index_path)
        self.node_features = self.code_index_map.set_index('code')[['node_x', 'node_y', 'map_code']]
        self.code_index_map = self.code_index_map.set_index('code')

    def get_node_index_by_code(self, node_code):
        """通过节点代码获取索引"""
        try:
            return self.code_index_map.at[node_code, 'index']
        except KeyError:
            print(f"Node does not exist: {node_code}")
            # 返回一个默认的节点索引，例如最后一个节点的索引
            return self.code_index_map.index[-1]  # 或者返回 None 或其他默认处理

    def create_graph_from_files(self):
        """
        创建邻接矩阵以及节点和边的特征字典
        :return:
        """
        # 读取节点索引映射和类型信息
        code_index_map = pd.read_excel(self.code_index_path)
        node_indices = pd.Series(code_index_map['index'].values, index=code_index_map['code']).to_dict()
        self.node_features = code_index_map.set_index('index')[['node_x', 'node_y', 'map_code', 'is_lift']].to_dict('index')
        self.code_index_map = code_index_map.set_index('code')

        # 初始化节点数和邻接矩阵
        nodes_num = len(node_indices)
        self.adjacency_matrix = np.full((nodes_num, nodes_num), np.inf)  # 使用np.inf表示无连接

        # 初始化边特征字典
        self.edge_features = {}

        # 读取电梯与普通点的边信息，将电梯与普通点的路径长度设置为固定的40(取历史电梯任务平均耗时80/2)
        node_lifts = self.database_reader.read_robot_node_lift()
        for _, edge in node_lifts.iterrows():
            start_node_index = node_indices[edge['node_code']]
            end_node_index = node_indices[edge['lift_code']]
            self.adjacency_matrix[start_node_index, end_node_index] = 40
            self.adjacency_matrix[end_node_index, start_node_index] = 40
            self.edge_features[(start_node_index, end_node_index)] = {
                'edge_length': 40,
                'robot_velocity': 1,
                'taketime': 40
            }
            self.edge_features[(end_node_index, start_node_index)] = {
                'edge_length': 40,
                'robot_velocity': 1,
                'taketime': 40
            }

        # 读取普通边信息
        edges = self.database_reader.read_robot_edge()
        for _, edge in edges.iterrows():
            start_node_index = node_indices[edge['start_node']]
            end_node_index = node_indices[edge['end_node']]

            # 设置邻接矩阵
            self.adjacency_matrix[start_node_index, end_node_index] = edge['edge_length']

            # 添加边特征
            self.edge_features[(start_node_index, end_node_index)] = {
                'edge_length': edge['edge_length'],
                'robot_velocity': edge['robot_velocity'],
                'taketime': edge['taketime']
            }

            # if edge['type_name'] == '双向直线':
            #     self.adjacency_matrix[end_node_index, start_node_index] = edge['edge_length']
            #     self.edge_features[(end_node_index, start_node_index)] = {
            #         'edge_length': edge['edge_length'],
            #         'robot_velocity': edge['robot_velocity'],
            #         'taketime': edge['taketime']
            #     }
            self.adjacency_matrix[end_node_index, start_node_index] = edge['edge_length']
            self.edge_features[(end_node_index, start_node_index)] = {
                'edge_length': edge['edge_length'],
                'robot_velocity': edge['robot_velocity'],
                'taketime': edge['taketime']
            }

        # 构建图
        self.graph = nx.Graph()
        for i in range(self.adjacency_matrix.shape[0]):
            for j in range(i + 1, self.adjacency_matrix.shape[1]):
                if self.adjacency_matrix[i, j] != 0 and self.adjacency_matrix[i, j] != np.inf:
                    self.graph.add_edge(i, j, weight=self.adjacency_matrix[i, j])


    def plot_map_graph(self, map_code=21):
        """
        绘制楼层地图
        :param map_code: 楼层code
        :return:
        """
        if self.node_features is None or self.adjacency_matrix is None:
            print("Graph data has not been initialized. Please call create_graph_from_files() first.")
            return

        # 构建图
        G = nx.Graph()
        for index, feature in self.node_features.items():
            if feature['map_code'] != map_code:
                continue
            G.add_node(index, pos=(feature['node_x'], feature['node_y']))
            for i in range(index, self.adjacency_matrix.shape[0]):
                if self.node_features[i]['map_code'] == map_code and \
                        (self.adjacency_matrix[index][i] != np.inf or self.adjacency_matrix[i][index] != np.inf):
                    G.add_edge(index, i)

        # 绘制散点图
        plt.figure(figsize=(8, 6))
        pos = nx.get_node_attributes(G, 'pos')
        nx.draw_networkx_nodes(G, pos, node_size=20, node_color='skyblue', alpha=0.8)
        nx.draw_networkx_edges(G, pos, width=1.0, alpha=0.5, edge_color='gray')

        plt.title('Node Scatter Plot with Connections')
        plt.xlabel('node_x')
        plt.ylabel('node_y')
        plt.axis('equal')
        plt.show()

    def find_shortest_path(self, start_index, end_index):
        """
        给定两个节点找到最短路径，用来生成模拟的历史行驶路径or预测的未来路径
        :param start_index: 起点index
        :param end_index: 终点index
        :return: 最短路径和路径长度
        """
        if self.adjacency_matrix is None:
            print("Graph data has not been initialized. Please call create_graph_from_files() first.")
            return

        # # 构建图
        # G = nx.Graph()
        # for i in range(self.adjacency_matrix.shape[0]):
        #     for j in range(i + 1, self.adjacency_matrix.shape[1]):
        #         if self.adjacency_matrix[i, j] != 0 and self.adjacency_matrix[i, j] != np.inf:
        #             G.add_edge(i, j, weight=self.adjacency_matrix[i, j])

        # 寻找最短路径
        try:
            path = nx.shortest_path(self.graph, source=start_index, target=end_index, weight='weight')
            path_length = nx.shortest_path_length(self.graph, source=start_index, target=end_index, weight='weight')
            return path, path_length
        except nx.NetworkXNoPath:
            print(f"No path exists between {start_index} and {end_index}")
        except nx.NodeNotFound:
            print(f"One of the nodes {start_index} or {end_index} does not exist in the graph")
        return None, -1


    def get_node_index_by_code(self, node_code):
        """
        通过节点code得到index
        :param node_code: code
        :return:
        """
        try:
            return self.code_index_map.at[node_code, 'index']
        except KeyError as e:
            print(f"node doesn't exist: {e}")


    def get_edge_features(self, start_index, end_index):
        """
        给定两个节点索引，获取边特征
        :param start_index: 起点index
        :param end_index: 终点idnex
        :return: 边的特征
        """
        if self.edge_features is None:
            print("Edge features have not been initialized.")
            return None
        try:
            return self.edge_features[(start_index, end_index)]
        except KeyError as e:
            print(f"edge doesn't exist: {e}")
            return None

    def get_adj_matrix(self):
        return self.adjacency_matrix

    def get_node_code_by_index(self, index):
        try:
            return self.code_index_map.index[self.code_index_map['index'] == index].tolist()[0]
        except IndexError:
            print(f"Index does not exist: {index}")
            return None



if __name__ == '__main__':
    # 创建图
    handler = MapLoader()
    handler.create_graph_from_files()
    print(handler.get_node_code_by_index(405))
    print(handler.node_features)
    handler.plot_map_graph(map_code=46)
    # path, path_length = handler.find_shortest_path(0, 753)
    # print(path, path_length)
    # print(handler.get_node_index_by_code('021$SITE-00393'))
